namespace NBody.Analysis
{
    using System;
    using System.IO;
    
    public class Energy
    {
        public static void CalculatePotentialDirect(ref NBodySystem s, double eps)
        {
            int i, j;
            double x, y, z, r;
            double eps2 = eps * eps;
            
            Console.Error.WriteLine("Calculating the potential of each particle (direct method)");
            
            // check first particle to see if it's an EnergyParticle ...
            if (s[0].GetType() != typeof(NBody.EnergyParticle))
            {
                for (i = 0; i < s.NumParts; i++)
                    s[i] = new EnergyParticle(s[i]);
            }
            
            // make sure potential is zeroed
            for (i = 0; i < s.NumParts; i++)
                ((EnergyParticle)s[i]).Potential = 0.0f;
            
            // this is the long way, obviously ...
            for (i = 0; i < s.NumParts; i++)
            {
                for (j = i+1; j < s.NumParts; j++)
                {
                    x = s[j].X - s[i].X;
                    y = s[j].Y - s[i].Y;
                    z = s[j].Z - s[i].Z;
                    r = Math.Sqrt(x * x + y * y + z * z + eps2);
                    ((EnergyParticle)s[i]).Potential -= (double)(s[j].Mass / r);
                    ((EnergyParticle)s[j]).Potential -= (double)(s[i].Mass / r);
                }
            }
        }
        
        public static void CalculatePotentialShell(ref NBodySystem s, double eps)
        {
            Console.Error.WriteLine("Calculating the potential of each particle (shell method)");
            // check first particle to see if it's an EnergyParticle ...
            if (s[0].GetType() != typeof(NBody.EnergyParticle))
            {
                for (int i = 0; i < s.NumParts; i++)
                    s[i] = new EnergyParticle(s[i]);
            }
            
            // system must be sorted by radius
            if (!s.IsSortedByRadius)
                s.Sort();
            
            double eps2 = eps*eps;
            
            double out_shell = 0.0;
            double[] r = new double[s.NumParts];
            for (int i = 0; i < s.NumParts; i++)
            {
                double radius = s[i].Radius;
                r[i] = Math.Sqrt(radius * radius + eps2);
                out_shell += s[i].Mass / r[i];
            }
            
            double enc_mass = 0.0;
            for (int i = 0; i < s.NumParts; i++)
            {
                enc_mass += s[i].Mass;
                out_shell -= s[i].Mass / r[i];
                ((EnergyParticle)s[i]).Potential = (double)(-1.0 * enc_mass / r[i] - out_shell);
            }
        }
        
        public static void SetPotentialFromData(ref NBodySystem s, DataFile.Table data)
        {
            Console.Error.WriteLine("Setting the potential of each particle from file");
            // check first particle to see if it's an EnergyParticle ...
            if (s[0].GetType() != typeof(NBody.EnergyParticle))
            {
                for (int i = 0; i < s.NumParts; i++)
                    s[i] = new EnergyParticle(s[i]);
            }
            
            for (int i = 0; i < data.Length; i++)
                ((EnergyParticle)s[i]).Potential = (double)data.X(i);
        }
        
        public static void SetPotentialFromData(ref NBodySystem s, double[] phi)
        {
            Console.Error.WriteLine("Setting the potential of each particle from array");
            // check first particle to see if it's an EnergyParticle ...
            if (s[0].GetType() != typeof(NBody.EnergyParticle))
            {
                for (int i = 0; i < s.NumParts; i++)
                    s[i] = new EnergyParticle(s[i]);
            }
            
            for (int i = 0; i < phi.Length; i++)
                ((EnergyParticle)s[i]).Potential = (double)phi[i];
        }
        
        public static void CalculateTotalEnergy(ref NBodySystem s, out double total_energy)
        {
            Console.Error.WriteLine("Calculating the energy of each particle");
            
            // check first particle to see if it's an EnergyParticle ...
            if (s[0].GetType() != typeof(NBody.EnergyParticle))
                throw new NBodyException("CalculateTotalEnergy: Particles do not contain potential information");
            
            total_energy = 0.0;
            for (int i = 0; i < s.NumParts; i++)
            {
                double ke = s[i].AbsoluteVelocity;
                ke *= 0.5 * ke;
                EnergyParticle p = s[i] as EnergyParticle;
                p.Energy = p.Potential + (double)ke;
                
                total_energy += p.Energy;
            }
        }
    }
}            
